import copy

import torch
import torch.nn.functional as F
from loguru import logger
from torch import nn
from torch.nn.init import xavier_uniform_

is_print_shape = False


class MyTransformer(nn.Module):
    def __init__(self, d_model=512,
                 nhead=8,
                 num_encoder_layers=6,
                 num_decoder_layers=6,
                 dim_feedforward=2048,
                 dropout=0.1):
        """
        :param d_model: d_k = d_v = d_model/nhead=64, 模型中向量的维度，默认值为 512
        :param nhead: 多头注意力机制中多头的数量，论文默认为值 8
        :param num_encoder_layers: encoder 堆叠的数量，论文中的 N，论文默认值为 6
        :param num_decoder_layers: decoder 堆叠的数量，论文中的 N，论文默认值为 6
        :param dim_feedforward: 全连接中向量的维度，论文默认值为 2048
        :param dropout: 丢弃率，论文中的默认值为 0.1
        """
        super(MyTransformer, self).__init__()
        # ================ 编码部分 =====================
        encoder_layer = MyTransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout)
        encoder_nor = nn.LayerNorm(d_model)
        self.encoder = MyTransformerEncoder(encoder_layer, num_encoder_layers, encoder_nor)
        # ================ 解码部分 =====================
        decoder_layer = MyTransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout)
        decoder_norm = nn.LayerNorm(d_model)

        self.decoder = MyTransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)

        # 初始化模型参数
        self._reset_parameters()

        self.d_model = d_model
        self.nhead = nhead

    def forward(self, src, tgt,
                src_mask=None,
                tgt_mask=None,
                memory_mask=None,
                src_key_padding_mask=None,
                tgt_key_padding_mask=None,
                memory_key_padding_mask=None):
        """

        :param src: [src_len,batch_size,embed_dim]
        :param tgt: [tgt_len, batch_size, embed_dim]
        :param src_mask: None
        :param tgt_mask: [tgt_len, tgt_len]
        :param memory_mask: None
        :param src_key_padding_mask:  [batch_size, src_len]
        :param tgt_key_padding_mask:  [batch_size, tgt_len]
        :param memory_key_padding_mask: [batch_size, src_len]
        :return:
                [tgt_len, batch_size, num_heads * kdim]
        """
        # [src_len, batch_size, num_heads * kdim]
        memory = self.encoder(src, mask=src_mask,
                              src_key_padding_mask=src_key_padding_mask)

        output = self.decoder(tgt=tgt, memory=memory,
                              tgt_mask=tgt_mask,
                              memory_mask=memory_mask,
                              tgt_key_padding_mask=tgt_key_padding_mask,
                              memory_key_padding_mask=memory_key_padding_mask)

        return output

    def generate_square_subsequent_mask(self, sz):
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))

        # mask shape [sz, sz]
        return mask

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                xavier_uniform_(p)


def _get_clones(modlue, N):
    return nn.ModuleList([copy.deepcopy(modlue) for _ in range(N)])


class MyTransformerDecoder(nn.Module):
    def __init__(self, decoder_layer, num_layers, norm=None):
        super(MyTransformerDecoder, self).__init__()
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, tgt, memory, tgt_mask=None,
                memory_mask=None,
                tgt_key_padding_mask=None,
                memory_key_padding_mask=None):
        # [tgt_len, batch_size, embed_dim]
        output = tgt

        # layers 为 N 层解码层堆叠起来
        for mod in self.layers:
            output = mod(output, memory,
                         tgt_mask=tgt_mask,
                         memory_mask=memory_mask,
                         tgt_key_padding_mask=tgt_key_padding_mask,
                         memory_key_padding_mask=memory_key_padding_mask)
        if self.norm is not None:
            output = self.norm(output)

        return output


class MyTransformerDecoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        """

        :param d_model: d_k = d_v = d_model/nhead= 64,模型中向量的维度，默认值为 512
        :param n_head: 多头注意力机制中多头的数量，论文默认为值 8
        :param dim_feedforward: 全连接中向量的维度，论文默认值为 2048
        :param dropout: 丢弃率，论文中的默认值为 0.1
        """
        super(MyTransformerDecoderLayer, self).__init__()

        # 解码部分输入序列之间的多头注意力（即论文结构图中的 Masked Multi-head attention)
        self.self_attn = MyMultiheadAttention(embed_dim=d_model, num_heads=nhead, dropout=dropout)
        # 编码部分输出（memory）和解码部分之间的多头注意力机制
        self.multihead_attn = MyMultiheadAttention(embed_dim=d_model, num_heads=nhead, dropout=dropout)

        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
        self.activation = nn.ReLU()

    def forward(self, tgt,
                memory,
                tgt_mask=None,
                memory_mask=None,
                tgt_key_padding_mask=None,
                memory_key_padding_mask=None):
        # 解码部分输入序列之间的多头注意力（论文结构图中的 Masked Multi-head attention)
        tgt2 = self.self_attn(tgt, tgt, tgt,  # [tgt_len,batch_size,embed_dim]
                              attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        # 残差连接
        tgt = tgt + self.dropout1(tgt2)
        # [tgt_len, batch_size, embed_dim]
        tgt = self.norm1(tgt)

        # 解码部分的输入经多头注意力后同编码部分的输出（memory）通过多头注意力机制进行交互
        # [tgt_len, batch_size, embed_dim]
        tgt2 = self.multihead_attn(tgt, memory, memory,
                                   attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        # [tgt_len, batch_size, dim_feedforward]
        tgt2 = self.activation(self.linear1(tgt))
        # [tgt_len, batch_size, embed_dim]
        tgt2 = self.linear2(self.dropout2(tgt2))

        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)

        return tgt


class MyTransformerEncoder(nn.Module):
    def __init__(self, encoder_layer, num_layers, norm=None):
        """
        :param encoder_layer: 包含有多头注意力机制的一个编码层
        :param num_layers: 克隆得到多个 encoder layers 论文中默认为 6
        :param norm: 归一化层
        """
        super(MyTransformerEncoder, self).__init__()
        # 克隆得到多个 encoder layers 论文中默认为 6
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, src, mask=None, src_key_padding_mask=None):
        """
        :param src: 编码部分的输入，形状为 [src_len,batch_size, embed_dim]
        :param mask: 编码部分输入的 padding 情况，形状为 [batch_size, src_len]
        :param src_key_padding_mask: None
        :return:
            Transformer Encoder编码器输出，[src_len, batch_size, num_heads * kdim] <==> [src_len, batch_size, embed_dim]
        """
        output = src
        for mod in self.layers:
            # 多个 encoder layers 层堆叠后的前向传播过程
            output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
        if self.norm is not None:
            output = self.norm(output)

        # [src_len, batch_size, num_heads * kdim] <==> [src_len, batch_size, embed_dim]
        return output


class MyTransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        """
        :param d_model: d_k = d_v = d_model/nhead=64,模型中向量的维度，默认值为 512
        :param nhead: 多头注意力机制中多头的数量，论文默认为值 8
        :param dim_feedforward: 全连接中向量的维度，论文默认值为 2048
        :param dropout: 丢弃率，论文中的默认值为 0.1
        """
        super(MyTransformerEncoderLayer, self).__init__()
        self.self_attn = MyMultiheadAttention(d_model, nhead, dropout=dropout)
        self.dropout1 = nn.Dropout(dropout)
        self.norm1 = nn.LayerNorm(d_model)

        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.activation = nn.ReLU()
        self.dropout2 = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        """

        :param src: 编码部分的输入，形状为 [src_len, batch_size, embed_dim]
        :param src_mask: 编码部分输入的padding情况，形状为 [batch_size, src_len]
        :param src_key_padding_mask: None
        :return:
            [src_len, batch_size, num_heads * kdim] <=> [src_len, batch_size, embed_dim]
        """

        # 计算多头注意力
        # src2 shape = [src_len, batch_size, num_heads * kdim]
        # num_heads * kdim = embed_dim
        src2 = self.self_attn(src, src, src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0]
        # 残差连接
        src = src + self.dropout1(src2)
        # src shape = [src_len, batch_size, num_heads * kdim]
        src = self.norm1(src)

        # src2 shape = [src_len, batch_size, dim_feedforward]
        src2 = self.activation(self.linear1(src))
        # [src_len, batch_size, num_heads * kdim]
        src2 = self.linear2(self.dropout(src2))
        src = src + self.dropout2(src2)
        src = self.norm2(src)

        # [src_len, batch_size, num_heads * kdim] <=> [src_len, batch_size, embed_dim]
        return src


class MyMultiheadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0., bias=True):
        super(MyMultiheadAttention, self).__init__()
        """
        :param embed_dim: 词嵌入的维度，也就是前面的 d_model 参数,论文中的默认值为 512
        :param num_heads: 多头注意力中多头的数量,也就是前面的 nhead 参数,论文默认值为 8
        :param bias: 最后对多头的注意力（组合）输出进行线性变换时，是否使用偏置
        """
        self.embed_dim = embed_dim
        self.head_dim = embed_dim // num_heads
        self.kdim = self.head_dim
        self.vdim = self.head_dim
        self.num_heads = num_heads  # 多头数量
        self.dropout = dropout

        # 以下限制条件即论文中的 d_k = d_v = d_model/n_head 条件
        assert self.head_dim * num_heads == self.embed_dim, "embed_dim 除以 num_heads 必须为整数"

        # embed_dim = kdim * num_heads
        # 这里第二个维度之所以是 embed_dim，实际上这里是同时初始化了 num_heads 个 W_q 堆叠起来的, 也就是 num_heads 个头
        self.q_proj_weight = nn.Linear(embed_dim, embed_dim, bias=bias)
        # W_k, embed_dim = kdim * num_heads
        self.k_proj_weight = nn.Linear(embed_dim, embed_dim, bias=bias)
        # W_v, embed_dim = vdim * num_heads
        self.v_proj_weight = nn.Linear(embed_dim, embed_dim, bias=bias)
        # 最后将所有的 Z 组合起来的时候，也是一次性完成， embed_dim = vdim * num_heads
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)

        self._reset_parameters()

    def _reset_parameters(self):
        """
        以特定的方式初始化参数
        """
        for p in self.parameters():
            if p.dim() > 1:
                xavier_uniform_(p)

    def forward(self, query, key, value, attn_mask=None, key_padding_mask=None):
        """
        在论文中，编码时 query, key, value 都是同一个输入，
        解码时 输入的部分也都是同一个输入，
        解码和编码交互时 key,value 指的是 memory, query 指的是 tgt

        :param query: [tgt_len, batch_size, embed_dim], tgt_len 表示目标序列的长度
        :param key: [src_len, batch_size, embed_dim], src_len 表示源序列的长度
        :param value: [src_len, batch_size, embed_dim], src_len 表示源序列的长度
        :param attn_mask:[tgt_len, src_len] or [num_heads * batch_size, tgt_len, src_len]
                一般只在解码时使用，为了并行一次喂入所有解码部分的输入，所以要用 mask 来进行掩盖当前时刻之后的位置信息
        :param key_padding_mask: [batch_size, src_len], src_len 表示源序列的长度

        :return:
            attn_output: [tgt_len, batch_size, embed_dim]
            attn_output_weights: # [batch_size, tgt_len, src_len]
        """
        return multi_head_attention_forward(query, key, value,
                                            self.num_heads,
                                            self.dropout,
                                            out_proj_weight=self.out_proj,
                                            training=self.training,
                                            key_padding_mask=key_padding_mask,
                                            q_proj_weight=self.q_proj_weight,
                                            k_proj_weight=self.k_proj_weight,
                                            v_proj_weight=self.v_proj_weight,
                                            attn_mask=attn_mask)


def multi_head_attention_forward(query,  # [tgt_len,batch_size, embed_dim]
                                 key,  # [src_len, batch_size, embed_dim]
                                 value,  # [src_len, batch_size, embed_dim]
                                 num_heads,
                                 dropout_p,
                                 out_proj_weight,  # [embde_dim = vdim * num_heads, embed_dim]
                                 training=True,
                                 key_padding_mask=None,  # [batch_size, src_len/ tgt_len]
                                 q_proj_weight=None,  # [embed_dim, kdim * num_heads]
                                 k_proj_weight=None,  # [embed_dim, kdim * num_heads]
                                 v_proj_weight=None,  # [embed_dim, vdim * num_heads]
                                 attn_mask=None  # [tgt_len, src_len]
                                 ):
    # 第一阶段：计算得到Q、K、V
    # [tgt_len, batch_size, embed_dim] X [embed_dim, kdim * num_heads] = [tgt_len, batch_size, kdim * num_heads]
    q = q_proj_weight(query)

    # [src_len, batch_size, embed_dim] X [embed_dim, kdim * num_heads] = [src_len, batch_size, kdim * num_heads]
    k = k_proj_weight(key)

    # [src_len, batch_size, embed_dim] X [embed_dim, vdim * num_heads] = [src_len, batch_size, vdim * num_heads]
    v = v_proj_weight(value)
    if is_print_shape:
        logger.info("=" * 40 + " Transformer Multihead Attention Caculate Start " + "=" * 40)
        logger.debug(f"num_heads = {num_heads}")
        logger.debug(f"d_model = {query.size(-1)}")
        logger.debug(f"d_k = d_v = d_model/num_heads = {query.size(-1) // num_heads}")
        logger.debug(f"query.shape = ([tgt_len, batch_size, embed_dim]) = {query.shape}")
        logger.debug(f"W_q.shape = ([embed_dim,kdim * num_heads]) = {q_proj_weight.shape}")
        logger.debug(f"Q.shape = ([tgt_len, batch_size,kdim * num_heads]) = {q.shape}")
        logger.debug("" + "-" * 70)

        logger.debug(f"key.shape = ([src_len,batch_size, embed_dim]) = {key.shape}")
        logger.debug(f"W_k.shape = ([embed_dim,kdim * num_heads]) = {k_proj_weight.shape}")
        logger.debug(f"K.shape = ([src_len,batch_size,kdim * num_heads]) = {k.shape}")
        logger.debug("" + "-" * 70)

        logger.debug(f"value.shape = ([src_len,batch_size, embed_dim]) = {value.shape}")
        logger.debug(f"W_v.shape = ([embed_dim,vdim * num_heads]) = {v_proj_weight.shape}")
        logger.debug(f"V.shape = ([src_len,batch_size,vdim * num_heads]) = {v.shape}")
        logger.debug("W_q,W_k,W_v同时进行计算。因此，Q,K,V分别也是包含了多个head的q,k,v堆叠起来的结果")
        logger.info("=" * 40 + " Transformer Multihead Attention Caculate Finish " + "=" * 40)

    # 第二阶段：缩放，以及attn_mask维度判断
    tgt_len, bsz, embed_dim = query.size()  # [tgt_len,batch_size,embed_dim]
    src_len = key.size(0)
    head_dim = embed_dim // num_heads  # num_heads * head_dim = embed_dim
    scaling = float(head_dim) ** -0.5
    q = q * scaling  # [query_len, batch_size, kdim * num_heads]

    if attn_mask is not None:
        # [tgt_len, src_len] or [num_heads * batch_size, tgt_len, src_len]
        if attn_mask.dim() == 2:
            attn_mask = attn_mask.unsqueeze(0)  # [1, tgt_len, src_len] 扩充维度
            if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
                raise RuntimeError('The size of 2D attn_mask is not correct.')
        elif attn_mask.dim() == 3:
            if list(attn_mask.size()) != [bsz * num_heads, query.size(0), key.size(0)]:
                raise RuntimeError('The size of the 3D attn_mask is not correct.')

    # 第三阶段： 计算得到注意力权重矩阵

    # 因为前面是 num_heads 个头一起参与的计算，所以这里要进行一下变形，以便于后面计算。且同时交换了 0，1 两个维度
    # q shape = # [batch_size * num_heads,tgt_len,kdim]
    q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
    # k shape = [batch_size * num_heads,src_len,kdim]
    k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
    # v shape = [batch_size * num_heads,src_len,vdim]
    v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)

    # [batch_size * num_heads,tgt_len,kdim] x [batch_size * num_heads, kdim, src_len] = [batch_size*num_heads, tgt_len, src_len]
    # num_heads 个 QK 相乘后的注意力矩阵
    attn_output_weights = torch.bmm(q, k.transpose(1, 2))

    # 第四阶段：进行相关掩码操作
    if attn_mask is not None:
        attn_output_weights += attn_mask  # [batch_size * num_heads, tgt_len, src_len]
    if key_padding_mask is not None:
        # 维度转换为 [bsz, num_heads, tgt_len, src_len]
        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
        attn_output_weights = attn_output_weights.masked_fill(key_padding_mask.unsqueeze(1).unsqueeze(2),
                                                              float('-inf'))
        # 维度还原为 [bsz * num_heads, tgt_len, src_len]
        attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len, src_len)

    # [batch_size * num_heads, tgt_len, src_len]
    attn_output_weights = F.softmax(attn_output_weights, dim=-1)

    attn_output_weights = F.dropout(attn_output_weights, p=dropout_p, training=training)

    # Z = [batch_size * num_heads, tgt_len, src_len] @ [batch_size * num_heads, src_len, vdim]
    # num_heads 个 Attention(Q, K, V) 结果
    attn_output = torch.bmm(attn_output_weights, v)

    # 首先 transpose 成 [tgt_len, batch_size * num_heads, kdim]
    # 然后 view 成 [tgt_len, batch_size, num_heads * kdim]
    attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)

    attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)

    # 这里就是多个 z 线性组合成 Z [tgt_len, batch_size, embed_dim]
    Z = out_proj_weight(attn_output)

    return Z, attn_output_weights.sum(dim=1) / num_heads  # 将 num_heads 个注意力权重矩阵按对应维度取平均
